package org.yu.api.register.data;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import net.minecraft.core.Registry;
import net.minecraft.data.loot.BlockLoot;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.storage.loot.BuiltInLootTables;
import net.minecraft.world.level.storage.loot.LootTable;
import org.jetbrains.annotations.NotNull;
import org.yu.api.register.BlockHolder;


import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;

public class LazyBlockLoot extends BlockLoot {
    private final Map<ResourceLocation, LootTable.Builder> map = Maps.newHashMap();
    private final Map<BlockHolder<?>, BiConsumer<BlockHolder<?>, LazyBlockLoot>> blockHolderConsumerMap = new HashMap<>();

    @Override
    protected void addTables() {
        blockHolderConsumerMap.forEach((blockHolder, lazyBlockLootConsumer) -> {
            lazyBlockLootConsumer.accept(blockHolder, this);
        });
    }

    public void put(BlockHolder<?> blockHolder, BiConsumer<BlockHolder<?>, LazyBlockLoot> loot) {
        this.blockHolderConsumerMap.put(blockHolder, loot);
    }

    @Override
    public void accept(@NotNull BiConsumer<ResourceLocation, LootTable.Builder> builderBiConsumer) {
        this.addTables();
        Set<ResourceLocation> set = Sets.newHashSet();
        for (Block block : Registry.BLOCK) {
            ResourceLocation resourcelocation = block.getLootTable();
            if (resourcelocation != BuiltInLootTables.EMPTY && set.add(resourcelocation)) {
                LootTable.Builder builder = this.map.remove(resourcelocation);
                if (builder != null) {
                    builderBiConsumer.accept(resourcelocation, builder);
                }
            }
        }
    }

    @Override
    protected void add(@NotNull Block block, @NotNull LootTable.Builder builder) {
        this.map.put(block.getLootTable(), builder);
    }
}
